
(ns ken.pu.lucene-helper
  (:import 
    (java.io File StringReader)
    (org.apache.lucene.document Document Field Field$Store Field$Index)
    (org.apache.lucene.index Term IndexReader IndexWriter IndexWriter$MaxFieldLength)
    (org.apache.lucene.queryParser QueryParser)
    (org.apache.lucene.search IndexSearcher)
    (org.apache.lucene.store SimpleFSDirectory)
    (org.apache.lucene.analysis SimpleAnalyzer WhitespaceAnalyzer Token)
    (org.apache.lucene.analysis.tokenattributes OffsetAttribute TermAttribute)
    (org.apache.lucene.search BooleanQuery TermQuery BooleanClause$Occur)
    (org.apache.lucene.util Version))
  (:use [clojure.contrib.str-utils])
  (:require [clojure.string]))

(def version "0.1")

(defn mk-analyzer
  "Constructs an analyzer"
  [& args]
  (if (= :simple (nth args 1))
    (SimpleAnalyzer.)
    (WhitespaceAnalyzer.)))

(defn analyze-string
  "Apply an analyzer to a string"
  [analyzer string]
  (let [token-stream (. analyzer tokenStream nil (StringReader. string))
        termAttribute (.getAttribute token-stream TermAttribute)]
    (loop [result []]
      (if (.incrementToken token-stream)
        (recur (conj result (.term termAttribute)))
        result))))

(defn mk-index-writer
  "Constructs an index writer"
  [path]
  (let [dir (-> path File. SimpleFSDirectory.)]
    (IndexWriter. dir (mk-analyzer) IndexWriter$MaxFieldLength/UNLIMITED)))

(defn mk-index-searcher
  "Constructs an index searcher"
  [path]
  (IndexSearcher. (IndexReader/open (-> path File. SimpleFSDirectory.))))

(defn- mk-field
  "constructs a field with keys
  :f :v :no-store? :no-index? :no-analyze?"
  [field]
  (Field. (:f field) (str (:v field))
    (if (:no-store field) Field$Store/NO Field$Store/YES)
    (if (:no-index field) Field$Index/NO 
      (if (:no-analyze field) Field$Index/NOT_ANALYZED Field$Index/ANALYZED))))
 
(defn mk-doc
  "Constructs a lucene docucment (mk-doc [{:keys (f v no-store no-index)}])"
  [& fields]
  (let [doc (Document.)]
    (do
      (doseq [field fields] (. doc add (mk-field field)))
      doc)))

(defn index-doc 
  "adds doc to an index-writer"
  [idx-writer doc]
  (. idx-writer addDocument doc))

(defn close-index-writer
  "closes an index-writer"
  [idx-writer]
  (doto idx-writer (.commit) (.optimize) (.close)))

(defn entity-to-doc
  "converts a tuple row (a clojure struct) to a Lucene document.
   We expect :id and :__all__ in the row."
  [m]
  (apply mk-doc
    (for [[k v] m]
    (if (= k :id)
      {:f (name k) :v v :no-analyze true}
      {:f (name k) :v v}))))

(defn doc-to-hashmap [doc]
  (apply hash-map (flatten (for [field (. doc getFields)]
    [(keyword (.name field)) (.stringValue field)]))))

(defn mk-query
  "constructs a query
   (mk-query [{:f field :v value :and yes :not yes}])"
  [& query-defs]
  (let [q (BooleanQuery.)]
    (do
      (doseq [query-def query-defs]
        (. q add 
          (TermQuery. (Term. (:f query-def) (:v query-def)))
          (if (:and query-def)
            (BooleanClause$Occur/MUST)
            (if (:not query-def)
              (BooleanClause$Occur/MUST_NOT)
              (BooleanClause$Occur/SHOULD)))))
      q)))

(defn mk-simple-query
  "constructs a query by simply performing the analysis and join the results
   into SHOULD clauses"
  ([flag analyzer field string]
   (let [q (BooleanQuery.)
        clause (if (= flag :and) (BooleanClause$Occur/MUST) (BooleanClause$Occur/SHOULD))]
    (do
      (doseq [token (analyze-string analyzer string)]
        (. q add (TermQuery. (Term. field token)) clause))
      q)))
  ([flag analyzer string]
    (mk-simple-query flag analyzer "__all__" string))
  ([analyzer string]
    (mk-simple-query :or analyzer string)))


(defn index-search [index-searcher query topk]
  (let [topdocs (. index-searcher search query topk)]
    (for [scoredoc (. topdocs scoreDocs)]
      (let [doc (. index-searcher doc (. scoredoc doc))]
        (doc-to-hashmap doc)))))

(defn ngram
  ([string n padding-char]
    (if (nil? string)
      '()
      (let [padding (apply str (repeat (dec n) padding-char))
          s (str padding (clojure.string/replace string " " padding) padding)]
        (map #(subs s % (+ n %)) (range (- (count s) (dec n)))))))
  ([string n]
    (ngram string n "$"))
  ([string]
    (ngram string 3 "$")))

(defn ngram-query
  [& args]
    (let [q (BooleanQuery.)]
      (do (doseq [token (apply ngram args)]
            (. q add (TermQuery. (Term. "__all__" token)) BooleanClause$Occur/SHOULD))
        q)))

(defn value-to-doc
  [ent-type attr value]
  (mk-doc
    {:f "entity" :v (str ent-type) :no-index true}
    {:f "attr" :v (if (keyword? attr) (name attr) (str attr)) :no-index true}
    {:f "value" :v (str value) :no-index true}
    {:f "__all__" :v (str-join " " (ngram value))}))